#include "GrabCut.hpp"

using namespace cv;

double GrabCut::calcBeta(const Mat &img)
{
    double beta = 0;
    for (int y = 0; y < img.rows; y++)
    {
        for (int x = 0; x < img.cols; x++)
        {
            //计算四个方向邻域两像素的差别，也就是欧式距离或者说二阶范数
            //（当所有像素都算完后，就相当于计算八邻域的像素差了）
            Vec3d color = img.at<Vec3b>(y, x);
            if (x > 0) // left  >0的判断是为了避免在图像边界的时候还计算，导致越界
            {
                Vec3d diff = color - (Vec3d)img.at<Vec3b>(y, x - 1);
                beta += diff.dot(diff); //矩阵的点乘，也就是各个元素平方的和
            }
            if (y > 0 && x > 0) // upleft
            {
                Vec3d diff = color - (Vec3d)img.at<Vec3b>(y - 1, x - 1);
                beta += diff.dot(diff);
            }
            if (y > 0) // up
            {
                Vec3d diff = color - (Vec3d)img.at<Vec3b>(y - 1, x);
                beta += diff.dot(diff);
            }
            if (y > 0 && x < img.cols - 1) // upright
            {
                Vec3d diff = color - (Vec3d)img.at<Vec3b>(y - 1, x + 1);
                beta += diff.dot(diff);
            }
        }
    }
    if (beta <= std::numeric_limits<double>::epsilon())
        beta = 0;
    else
        beta = 1.f / (2 * beta / (4 * img.cols * img.rows - 3 * img.cols - 3 * img.rows + 2)); //论文公式（5）

    return beta;
}

void GrabCut::calcNWeights(const Mat &img, Mat &leftW, Mat &upleftW, Mat &upW,
                           Mat &uprightW, double beta, double gamma)
{
    //gammaDivSqrt2相当于公式（4）中的gamma * dis(i,j)^(-1)，那么可以知道，
    //当i和j是垂直或者水平关系时，dis(i,j)=1，当是对角关系时，dis(i,j)=sqrt(2.0f)。
    //具体计算时，看下面就明白了
    const double gammaDivSqrt2 = gamma / std::sqrt(2.0f);
    //每个方向的边的权值通过一个和图大小相等的Mat来保存
    leftW.create(img.rows, img.cols, CV_64FC1);
    upleftW.create(img.rows, img.cols, CV_64FC1);
    upW.create(img.rows, img.cols, CV_64FC1);
    uprightW.create(img.rows, img.cols, CV_64FC1);
    for (int y = 0; y < img.rows; y++)
    {
        for (int x = 0; x < img.cols; x++)
        {
            Vec3d color = img.at<Vec3b>(y, x);
            if (x - 1 >= 0) // left  //避免图的边界
            {
                Vec3d diff = color - (Vec3d)img.at<Vec3b>(y, x - 1);
                leftW.at<double>(y, x) = gamma * exp(-beta * diff.dot(diff));
            }
            else
                leftW.at<double>(y, x) = 0;
            if (x - 1 >= 0 && y - 1 >= 0) // upleft
            {
                Vec3d diff = color - (Vec3d)img.at<Vec3b>(y - 1, x - 1);
                upleftW.at<double>(y, x) = gammaDivSqrt2 * exp(-beta * diff.dot(diff));
            }
            else
                upleftW.at<double>(y, x) = 0;
            if (y - 1 >= 0) // up
            {
                Vec3d diff = color - (Vec3d)img.at<Vec3b>(y - 1, x);
                upW.at<double>(y, x) = gamma * exp(-beta * diff.dot(diff));
            }
            else
                upW.at<double>(y, x) = 0;
            if (x + 1 < img.cols && y - 1 >= 0) // upright
            {
                Vec3d diff = color - (Vec3d)img.at<Vec3b>(y - 1, x + 1);
                uprightW.at<double>(y, x) = gammaDivSqrt2 * exp(-beta * diff.dot(diff));
            }
            else
                uprightW.at<double>(y, x) = 0;
        }
    }
}

void GrabCut::checkMask(const Mat &img, const Mat &mask)
{
    if (mask.empty())
        CV_Error(Error::StsBadArg, "mask is empty");
    if (mask.type() != CV_8UC1)
        CV_Error(Error::StsBadArg, "mask must have CV_8UC1 type");
    if (mask.cols != img.cols || mask.rows != img.rows)
        CV_Error(Error::StsBadArg, "mask must have as many rows and cols as img");
    for (int y = 0; y < mask.rows; y++)
    {
        for (int x = 0; x < mask.cols; x++)
        {
            uchar val = mask.at<uchar>(y, x);
            if (val != GC_BGD && val != GC_FGD && val != GC_PR_BGD && val != GC_PR_FGD)
                CV_Error(Error::StsBadArg, "mask element value must be equel"
                                           "GC_BGD or GC_FGD or GC_PR_BGD or GC_PR_FGD");
        }
    }
}

void GrabCut::initMaskWithRect(Mat &mask, Size imgSize, Rect rect)
{
    mask.create(imgSize, CV_8UC1);
    mask.setTo(GC_BGD);

    rect.x = max(0, rect.x);
    rect.y = max(0, rect.y);
    rect.width = min(rect.width, imgSize.width - rect.x);
    rect.height = min(rect.height, imgSize.height - rect.y);

    (mask(rect)).setTo(Scalar(GC_PR_FGD));
    return;
}

void GrabCut::initGMMs(const Mat &img, const Mat &mask, GMM &bgdGMM, GMM &fgdGMM)
{
    const int kMeansItCount = 10;             //迭代次数
    const int kMeansType = KMEANS_PP_CENTERS; //Use kmeans++ center initialization by Arthur and Vassilvitskii

    Mat bgdLabels, fgdLabels;             //记录背景和前景的像素样本集中每个像素对应GMM的哪个高斯模型，论文中的kn
    vector<Vec3f> bgdSamples, fgdSamples; //背景和前景的像素样本集
    Point p;
    for (p.y = 0; p.y < img.rows; p.y++)
    {
        for (p.x = 0; p.x < img.cols; p.x++)
        {
            uchar pixel = mask.at<uchar>(p);
            //mask中标记为GC_BGD和GC_PR_BGD的像素都作为背景的样本像素
            if (mask.at<uchar>(p) == GC_BGD || mask.at<uchar>(p) == GC_PR_BGD)
                bgdSamples.push_back((Vec3f)img.at<Vec3b>(p));
            else // GC_FGD | GC_PR_FGD
                fgdSamples.push_back((Vec3f)img.at<Vec3b>(p));
        }
    }
    CV_Assert(!bgdSamples.empty() && !fgdSamples.empty());

    //kmeans中参数_bgdSamples为：每行一个样本
    //kmeans的输出为bgdLabels，里面保存的是输入样本集中每一个样本对应的类标签（样本聚为componentsCount类后）
    Mat _bgdSamples((int)bgdSamples.size(), 3, CV_32FC1, &bgdSamples[0][0]);
    kmeans(_bgdSamples, GMM::componentsCount, bgdLabels,
           TermCriteria(1, kMeansItCount, 0.0), 0, kMeansType);
    Mat _fgdSamples((int)fgdSamples.size(), 3, CV_32FC1, &fgdSamples[0][0]);
    kmeans(_fgdSamples, GMM::componentsCount, fgdLabels,
           TermCriteria(1, kMeansItCount, 0.0), 0, kMeansType);

    //经过上面的步骤后，每个像素所属的高斯模型就确定的了，那么就可以估计GMM中每个高斯模型的参数了。
    bgdGMM.initLearning();
    for (int i = 0; i < (int)bgdSamples.size(); i++)
        bgdGMM.addSample(bgdLabels.at<int>(i, 0), bgdSamples[i]);
    bgdGMM.endLearning();

    fgdGMM.initLearning();
    for (int i = 0; i < (int)fgdSamples.size(); i++)
        fgdGMM.addSample(fgdLabels.at<int>(i, 0), fgdSamples[i]);
    fgdGMM.endLearning();
}

void GrabCut::assignGMMsComponents(const Mat &img, const Mat &mask, const GMM &bgdGMM,
                                   const GMM &fgdGMM, Mat &compIdxs)
{
    Point p;
    for (p.y = 0; p.y < img.rows; p.y++)
    {
        for (p.x = 0; p.x < img.cols; p.x++)
        {
            Vec3d color = img.at<Vec3b>(p);
            //通过mask来判断该像素属于背景像素还是前景像素，再判断它属于前景或者背景GMM中的哪个高斯分量
            compIdxs.at<int>(p) = mask.at<uchar>(p) == GC_BGD || mask.at<uchar>(p) == GC_PR_BGD ? bgdGMM.whichComponent(color) : fgdGMM.whichComponent(color);
        }
    }
}

void GrabCut::learnGMMs(const Mat &img, const Mat &mask, const Mat &compIdxs, GMM &bgdGMM, GMM &fgdGMM)
{
    bgdGMM.initLearning();
    fgdGMM.initLearning();
    Point p;
    for (int ci = 0; ci < GMM::componentsCount; ci++)
    {
        for (p.y = 0; p.y < img.rows; p.y++)
        {
            for (p.x = 0; p.x < img.cols; p.x++)
            {
                if (compIdxs.at<int>(p) == ci)
                {
                    if (mask.at<uchar>(p) == GC_BGD || mask.at<uchar>(p) == GC_PR_BGD)
                        bgdGMM.addSample(ci, img.at<Vec3b>(p));
                    else
                        fgdGMM.addSample(ci, img.at<Vec3b>(p));
                }
            }
        }
    }
    bgdGMM.endLearning();
    fgdGMM.endLearning();
}

void GrabCut::constructGCGraph(const Mat &img, const Mat &mask, const GMM &bgdGMM, const GMM &fgdGMM, double lambda,
                               const Mat &leftW, const Mat &upleftW, const Mat &upW, const Mat &uprightW,
                               GCGraph &graph)
{
    int vtxCount = img.cols * img.rows;                                 //顶点数，每一个像素是一个顶点
    int edgeCount = 2 * (4 * vtxCount - 3 * (img.cols + img.rows) + 2); //边数，需要考虑图边界的边的缺失
                                                                        //通过顶点数和边数创建图。这些类型声明和函数定义请参考gcgraph.hpp
    graph.create(vtxCount, edgeCount);
    Point p;
    for (p.y = 0; p.y < img.rows; p.y++)
    {
        for (p.x = 0; p.x < img.cols; p.x++)
        {
            // add node
            int vtxIdx = graph.addVtx(); //返回这个顶点在图中的索引
            Vec3b color = img.at<Vec3b>(p);

            // set t-weights
            //计算每个顶点与Sink汇点t（代表背景）和源点Source（代表前景）连接的权值。
            //也即计算Gibbs能量（每一个像素点作为背景像素或者前景像素）的第一个能量项
            double fromSource, toSink;
            if (mask.at<uchar>(p) == GC_PR_BGD || mask.at<uchar>(p) == GC_PR_FGD)
            {
                //对每一个像素计算其作为背景像素或者前景像素的第一个能量项，作为分别与t和s点的连接权值
                fromSource = -log(bgdGMM(color));
                toSink = -log(fgdGMM(color));
            }
            else if (mask.at<uchar>(p) == GC_BGD)
            {
                //对于确定为背景的像素点，它与Source点（前景）的连接为0，与Sink点的连接为lambda
                fromSource = 0;
                toSink = lambda;
            }
            else // GC_FGD
            {
                fromSource = lambda;
                toSink = 0;
            }
            //设置该顶点vtxIdx分别与Source点和Sink点的连接权值
            graph.addTermWeights(vtxIdx, fromSource, toSink);

            // set n-weights  n-links
            //计算两个邻域顶点之间连接的权值。
            //也即计算Gibbs能量的第二个能量项（平滑项）
            if (p.x > 0)
            {
                double w = leftW.at<double>(p);
                graph.addEdges(vtxIdx, vtxIdx - 1, w, w);
            }
            if (p.x > 0 && p.y > 0)
            {
                double w = upleftW.at<double>(p);
                graph.addEdges(vtxIdx, vtxIdx - img.cols - 1, w, w);
            }
            if (p.y > 0)
            {
                double w = upW.at<double>(p);
                graph.addEdges(vtxIdx, vtxIdx - img.cols, w, w);
            }
            if (p.x < img.cols - 1 && p.y > 0)
            {
                double w = uprightW.at<double>(p);
                graph.addEdges(vtxIdx, vtxIdx - img.cols + 1, w, w);
            }
        }
    }
}

void GrabCut::estimateSegmentation(GCGraph &graph, Mat &mask)
{
    //通过最大流算法确定图的最小割，也即完成图像的分割
    graph.maxFlow();
    Point p;
    for (p.y = 0; p.y < mask.rows; p.y++)
    {
        for (p.x = 0; p.x < mask.cols; p.x++)
        {
            //通过图分割的结果来更新mask，即最后的图像分割结果。注意的是，永远都
            //不会更新用户指定为背景或者前景的像素
            if (mask.at<uchar>(p) == GC_PR_BGD || mask.at<uchar>(p) == GC_PR_FGD)
            {
                if (graph.inSourceSegment(p.y * mask.cols + p.x /*vertex index*/))
                    mask.at<uchar>(p) = GC_PR_FGD;
                else
                    mask.at<uchar>(p) = GC_PR_BGD;
            }
        }
    }
}

void GrabCut::operator()(Mat _img, Mat _mask, Rect rect,
                         Mat _bgdModel, Mat _fgdModel,
                         int iterCount, int mode)
{
    Mat img = _img;
    Mat &mask = _mask;
    Mat &bgdModel = _bgdModel;
    Mat &fgdModel = _fgdModel;

    if (img.empty())
        CV_Error(Error::StsBadArg, "image is empty");
    if (img.type() != CV_8UC3)
        CV_Error(Error::StsBadArg, "image mush have CV_8UC3 type");

    GMM bgdGMM(bgdModel), fgdGMM(fgdModel);
    Mat compIdxs(img.size(), CV_32SC1);

    if (mode == GC_INIT_WITH_RECT || mode == GC_INIT_WITH_MASK)
    {
        if (mode == GC_INIT_WITH_RECT)
            initMaskWithRect(mask, img.size(), rect);
        else // flag == GC_INIT_WITH_MASK
            checkMask(img, mask);
        initGMMs(img, mask, bgdGMM, fgdGMM);
    }

    if (iterCount <= 0)
        return;

    if (mode == GC_EVAL)
        checkMask(img, mask);

    const double gamma = 50;
    const double lambda = 9 * gamma;
    const double beta = calcBeta(img);

    Mat leftW, upleftW, upW, uprightW;
    calcNWeights(img, leftW, upleftW, upW, uprightW, beta, gamma);

    for (int i = 0; i < iterCount; i++)
    {
        GCGraph graph;
        assignGMMsComponents(img, mask, bgdGMM, fgdGMM, compIdxs);
        learnGMMs(img, mask, compIdxs, bgdGMM, fgdGMM);
        constructGCGraph(img, mask, bgdGMM, fgdGMM, lambda, leftW, upleftW, upW, uprightW, graph);
        estimateSegmentation(graph, mask);
    }
}


void changeBackGD(Mat& forGD, Mat& backGD, Mat& mask, Mat& dst, bool flag)
{
    assert(forGD.size() == mask.size());
    dst = Mat::zeros(backGD.size(), CV_8UC3);
    if (forGD.size() == backGD.size())
    {
        for (size_t i = 0; i < dst.rows; i++)
        {
            // 获取各个Mat的像素值
            uchar* forgd = forGD.ptr<uchar>(i);
            uchar* backgd = backGD.ptr<uchar>(i);
            uchar* maskpix = mask.ptr<uchar>(i);
            uchar* dstpix = dst.ptr<uchar>(i);
            for (size_t j = 0; j < dst.cols; j++)
            {
                // 当掩码为255时，赋予前景的像素值
                if (maskpix[j] == 0)
                {
                    for (int k = 0; k < 3; k++)
                    {
                        *dstpix++ = *backgd++;
                        forgd++;
                    }
                    
                }
                // 当掩码为0时，赋予背景的像素值
                else 
                {
                    for (int k = 0; k < 3; k++)
                    {
                        *dstpix++ = *forgd++;
                        backgd++;
                    }
                }
            }
        }
    }
    else
    {
        Mat rzForGD, rzMask;
        cv::resize(forGD, rzForGD, Size(forGD.size().width * backGD.size().height / forGD.size().height, backGD.size().height));
        cv::resize(mask, rzMask, Size(mask.size().width * backGD.size().height / mask.size().height, backGD.size().height));
        int interial = (backGD.size().width - rzForGD.size().width) / 2;

        for (size_t i = 0; i < dst.rows; i++)
        {
            // 获取各个Mat的像素值
            uchar* forgd = rzForGD.ptr<uchar>(i);
            uchar* backgd = backGD.ptr<uchar>(i);
            uchar* maskpix = rzMask.ptr<uchar>(i);
            uchar* dstpix = dst.ptr<uchar>(i);
            for (size_t j = 0; j < dst.cols; j++)
            {
                if (j < interial || dst.cols - j < interial)
                {
                    for (int k = 0; k < 3; k++)
                    {
                        if (flag)
                            *dstpix++ = *backgd++;
                        else
                        {
                            dstpix++;
                            backgd++;
                        } 
                    }
                }
                else
                {
                    auto mpix = maskpix[j - interial];
                    // 当掩码为0时，赋予背景的像素值
                    if (mpix == 0)
                    {
                        for (int k = 0; k < 3; k++)
                        {
                            if (flag)
                                *dstpix++ = *backgd++;
                            else
                            {
                                *dstpix++ = 0;
                                backgd++;
                            }
                            forgd++;
                        }
                        
                    }
                    // 当掩码为255时，赋予前景的像素值
                    else
                    {
                        for (int k = 0; k < 3; k++)
                        {
                            *dstpix++ = *forgd++;
                            backgd++;
                        }
                    }
                }
                
            }
        }
    }
}
